package com.tuniu.agents.common.agent;

import com.tuniu.agents.advisor.AgentChatMemoryAdvisor;
import com.tuniu.agents.agent.AgentManagerGroup;
import com.tuniu.agents.agent.AgentOptions;
import com.tuniu.agents.agent.OptionsAgent;
import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.common.handler.IStreamResponseHandler;
import com.tuniu.agents.message.MessageLogger;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

public class Recommender extends OptionsAgent {

    private MessageLogger messageLogger;
    protected IStreamResponseHandler streamResponseHandler ;

    /**
     * Recommender：
     * The user previously raised these requests:
     * 1: XXX
     * 2: XXX
     * The request of the current user is:
     * XXX
     * The current query result is:
     * XXX（MD）
     */
    public static final String MESSAGE_TEMPLATE = """
            {{user_message_his}}
            The request of the current user is:
            {{user_last_message}}
            The current query result is:
            {{user_message_current}}
            """;

    public Recommender(AgentManagerGroup agentManagerGroup, String agentId, String promptTemplateId) {
        this(agentManagerGroup, agentId, promptTemplateId, null,null);
    }

    public Recommender(AgentManagerGroup agentManagerGroup, String agentId, String promptTemplateId, MessageLogger messageLogger) {
        this(agentManagerGroup, agentId, promptTemplateId, messageLogger,null);
    }

    public Recommender(AgentManagerGroup agentManagerGroup, String agentId, String promptTemplateId, MessageLogger messageLogger,IStreamResponseHandler streamResponseHandler) {
        super(agentManagerGroup.getModelBuilder(),
                agentManagerGroup.getChatMemory(),
                agentManagerGroup.getAdvisors().stream().map(advisor -> advisor instanceof AgentChatMemoryAdvisor ? new AgentChatMemoryAdvisor(agentManagerGroup.getChatMemory(), false, true) : advisor).collect(Collectors.toList()),
                agentManagerGroup.getPromptManager(),
                agentManagerGroup.getChatEnvironmentMemory(),
                new AgentOptions(agentId, promptTemplateId));
        this.messageLogger = messageLogger;
        this.streamResponseHandler = streamResponseHandler;
    }

    public Recommender(AgentManagerGroup agentManagerGroup, String agentId, String promptTemplateId, MessageLogger messageLogger,IStreamResponseHandler streamResponseHandler,boolean saveInRecord,boolean saveOutRecord) {
        super(agentManagerGroup.getModelBuilder(),
                agentManagerGroup.getChatMemory(),
                agentManagerGroup.getAdvisors().stream().map(advisor -> advisor instanceof AgentChatMemoryAdvisor ? new AgentChatMemoryAdvisor(agentManagerGroup.getChatMemory(), saveInRecord, saveOutRecord) : advisor).collect(Collectors.toList()),
                agentManagerGroup.getPromptManager(),
                agentManagerGroup.getChatEnvironmentMemory(),
                new AgentOptions(agentId, promptTemplateId));
        this.messageLogger = messageLogger;
        this.streamResponseHandler = streamResponseHandler;
    }

    public Recommender(AgentManagerGroup agentManagerGroup, String agentId, String promptTemplateId, IStreamResponseHandler streamResponseHandler) {
        this(agentManagerGroup, agentId, promptTemplateId, null,streamResponseHandler);
    }

    public Recommender(AgentManagerGroup agentManagerGroup, String agentId, String promptTemplateId, IStreamResponseHandler streamResponseHandler,boolean saveInRecord,boolean saveOutRecord) {
        this(agentManagerGroup, agentId, promptTemplateId, null,streamResponseHandler,saveInRecord,saveOutRecord);
    }

    @Override
    public ChatClient.ChatClientRequestSpec chatClientRequestSpec(String conversationId, List<Message> messages, Map<String, Object> context) {
        return chatClient.prompt()
                .system(s -> s.params(chatEnvironmentMemory.get(conversationId)))
                .messages(messages)
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId).
                        param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, super.agentOptions.getAdvisorParam().get(CHAT_MEMORY_RETRIEVE_SIZE_KEY)).
                        param(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_TEMPLATE, Recommender.MESSAGE_TEMPLATE)
                        .param(AgentsConstants.AgentsAdvisorKey.ASSISTANT_MESSAGE_HISTORY_NEED, "false")
                        .param(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_HISTORY_NEED, "false")
                        .param(AgentChatMemoryAdvisor.CHAT_MEMORY_AGENT_ID_KEY, this.agentId));
    }
}
